# -*- coding:utf-8 -*-

import os
from typing import NoReturn, Literal

from src._typing import MatchLocation, OCRLocationListHandle
from src._base import logger
from src._items import OCR_MODEL_BASE_PATH
from src._items._ocr_reader import OCRReader

__all__ = ('EasyOCRReader',)


class EasyOCRReader(OCRReader):

    # 当前支持的OCR识别引擎
    _engine_type: str = 'easyocr'
    # 默认支持的语言
    _supported_languages: tuple = ('ch', 'en')
    # 默认支持的语言对应字典
    _language_map: dict = {'ch': 'ch_sim', 'en': 'en'}
    # 默认模型路径
    _model_path: str = os.path.abspath(os.path.join(OCR_MODEL_BASE_PATH, _engine_type))

    def __init__(self, image_path: str, language: str | list | None = None,
                 model_base_path: str | None = None) -> NoReturn:
        super().__init__(image_path, language, model_base_path)
        # 配置当前引擎
        self._engine = None
        # 处理支持的语言
        if language is None:
            language = ['ch_sim', 'en']
        self._language = self.__init_language(language)

    def __init_language(self, language: str | list):
        """
        初始化语言
        Args:
            language(str | list): 传参语言

        Returns:
            res_lang(list): 返回语言列表

        """
        if isinstance(language, str):
            if language not in self._supported_languages:
                raise ValueError(
                    f'传参语言（{language}）不在当前OCR引擎支持的范围内！（当前引擎支持的语言为：{self._supported_languages}）')
            return [self._language_map[language]]
        else:
            res_lang = []
            for lang in language:
                if lang not in self._supported_languages:
                    raise ValueError(
                        f'传参语言（{language}）不在当前OCR引擎支持的范围内！（当前引擎支持的语言为：{self._supported_languages}）')
                res_lang.append(self._language_map[lang])
            return res_lang

    def _init_engine(self) -> NoReturn:
        import easyocr
        self._engine = easyocr.Reader(self._language, model_storage_directory=self._model_path, download_enabled=False)

    def _get_result(self) -> list:
        """
        获取OCR识别结果
        Returns:
            res(list): OCR识别结果

        """
        if self._engine is None:
            self._init_engine()
        res = self._engine.readtext(self._image_path)
        self._ocr_full_result = res
        return res

    def __get_ocr_result_dict(self, ocr_full_result: list, is_filter_special_symbol: bool = False
                              ) -> dict[str, MatchLocation]:
        """
        获取OCR识别结果
        Args:
            ocr_full_result(list): OCR全部识别结果
            is_filter_special_symbol(bool): 返回的结果中是否过滤特殊字符

        Returns:
            res(dict[str, MatchLocation]): OCR识别结果

        """
        res = {}
        for item in ocr_full_result:
            loc_list = item[0]
            result_str = item[1]
            # 过滤特殊字符
            if is_filter_special_symbol:
                result_str = self._filter_special_chars(result_str)
            # 将OCR识别结果转换为OCRLocationListHandle对象
            ocr_h = OCRLocationListHandle(loc_list)
            # 将OCR识别结果添加到结果字典中（值为处理过的中间点坐标）
            res[result_str] = ocr_h.sort_and_get_centre_location()
        return res

    def __get_ocr_result_list(self, ocr_full_result: list, is_filter_special_symbol: bool = False) -> list[str]:
        """
        获取OCR识别结果
        Args:
            ocr_full_result(list): OCR全部识别结果
            is_filter_special_symbol(bool): 返回的结果中是否过滤特殊字符

        Returns:
            res(list[str]): OCR识别结果

        """
        res = []
        for item in ocr_full_result:
            result_str = item[1]
            # 过滤特殊字符
            if is_filter_special_symbol:
                result_str = self._filter_special_chars(result_str)
            res.append(result_str)
        return res

    def __get_ocr_result_str(self, ocr_full_result: list, is_filter_special_symbol: bool = False) -> str:
        """
        获取OCR识别结果
        Args:
            ocr_full_result(list): OCR全部识别结果
            is_filter_special_symbol(bool): 返回的结果中是否过滤特殊字符

        Returns:
            res(str): OCR识别结果

        """
        res = ''
        for item in ocr_full_result:
            result_str = item[1]
            # 过滤特殊字符
            if is_filter_special_symbol:
                result_str = self._filter_special_chars(result_str)
            res += result_str
        return res

    def get_picture_texts(self, result_tag: Literal['str', 'list', 'dict'] = 'str',
                          is_filter_special_symbol: bool = False
                          ) -> str | list[str] | dict[str, MatchLocation] | None:
        """
        获取OCR识别结果
        Args:
            result_tag(str): 返回结果的格式
            is_filter_special_symbol(bool): 返回的结果中是否过滤特殊字符

        Returns:
            result(str | list[str] | dict[str, MatchLocation] | None): OCR识别结果

        """
        if not self._ocr_full_result:
            self._get_result()
        match result_tag:
            case 'str':
                return self.__get_ocr_result_str(self._ocr_full_result, is_filter_special_symbol)
            case 'dict':
                return self.__get_ocr_result_dict(self._ocr_full_result, is_filter_special_symbol)
            case 'list':
                return self.__get_ocr_result_list(self._ocr_full_result, is_filter_special_symbol)
            case _:
                logger.error(f'当前返回结果的类型不支持：{result_tag}，请检查输入参数！')
                return None

    def text_is_match(self, text: str, is_filter_special_symbol: bool = True) -> bool:
        """
        文本匹配
        Args:
            text(str): 待匹配的文本
            is_filter_special_symbol(bool): 返回的结果中是否过滤特殊字符

        Returns:
            bool: 是否匹配成功

        """
        return text in self.get_picture_texts(is_filter_special_symbol=is_filter_special_symbol, result_tag='str')

    def texts_is_all_match(self, texts: list[str] | set[str] | tuple[str], is_filter_special_symbol: bool = True) -> bool:
        """
        文本匹配
        Args:
            texts(list[str] | set[str] | tuple[str]): 待匹配的文本
            is_filter_special_symbol(bool): 返回的结果中是否过滤特殊字符

        Returns:
            bool: 是否匹配成功

        """
        for text in texts:
            if not self.text_is_match(text, is_filter_special_symbol=is_filter_special_symbol):
                return False
        else:
            return True

    def text_match_location(self, text: str, is_filter_special_symbol: bool = True) -> MatchLocation | None:
        """
        获取文本在图片中的位置
        Args:
            text(str): 待匹配的文本
            is_filter_special_symbol(bool): 返回的结果中是否过滤特殊字符

        Returns:
            MatchLocation | None: 文本在图片中的位置

        """
        if not self.text_is_match(text, is_filter_special_symbol=is_filter_special_symbol):
            logger.error(f'文本：{text} 未在图片中匹配到！')
            return None
        text_match_dict = self.get_picture_texts(result_tag='dict', is_filter_special_symbol=is_filter_special_symbol)
        res = text_match_dict[text]
        return res

    def texts_match_location(self, texts: list[str] | set[str] | tuple[str],
                             is_filter_special_symbol: bool = True
                             ) -> list[MatchLocation] | list:
        """
        获取文本在图片中的位置
        Args:
            texts(list[str] | set[str] | tuple[str]): 待匹配的文本
            is_filter_special_symbol(bool): 返回的结果中是否过滤特殊字符

        Returns:
            list[MatchLocation] | list: 文本在图片中的位置。如果匹配失败，则返回空列表

        """
        res = []
        for text in texts:
            match_res = self.text_match_location(text, is_filter_special_symbol=is_filter_special_symbol)
            if match_res:
                res.append(match_res)
        return res


if __name__ == '__main__':

    import easyocr

    easy_ocr_model_path = os.path.join(OCR_MODEL_BASE_PATH, 'easyocr')

    e_ocr = easyocr.Reader(['ch_sim', 'en'], model_storage_directory=easy_ocr_model_path,
                           download_enabled=False)
    e_res = e_ocr.readtext('../test_pictures/numbers.jpeg')
    print(e_res)
